#include <opencv2/opencv.hpp>
#include <atomic>
#include <thread>
#include <mutex>
#include <vector>
#include <string>
#include <filesystem>
#include <numeric>
#include <iostream>
#include <cmath>
#include <chrono>
#include <iomanip>
#include <sstream>

#include "web_server.h"
#include "web_status.h"
#include "web_status_ui.h"
#include "web_utils.h"
#include "otl.h"

#include "RivsInference/RivsInferRuntime.h"

// -------------------- Simple logging utilities --------------------
namespace {
    inline std::string time_now_string() {
        using namespace std::chrono;
        auto now = system_clock::now();
        auto t_c = system_clock::to_time_t(now);
        auto ms = duration_cast<milliseconds>(now.time_since_epoch()) % 1000;
        std::tm tm{};
        localtime_r(&t_c, &tm);
        std::ostringstream oss;
        oss << std::put_time(&tm, "%H:%M:%S") << '.' << std::setw(3) << std::setfill('0') << ms.count();
        return oss.str();
    }
    static std::mutex g_log_mtx;
}

#define LOG(msg) do { std::lock_guard<std::mutex> _lk(g_log_mtx); \
    std::cerr << "[" << time_now_string() << "] [" << std::this_thread::get_id() << "] " << msg << std::endl; } while(0)

struct DetectorXConfig {
    std::string onnx_path;
    int device_id{0};
    std::vector<uint32_t> cluster_ids{0};
    int input_w{640};
    int input_h{640};
    float conf_thres{0.30f};
    float iou_thres{0.45f};
    bool p6{false};
    bool apply_sigmoid{false};
};

static inline float sigmoid(float x){ return 1.0f/(1.0f+std::exp(-x)); }

class YOLOXDetector {
public:
    explicit YOLOXDetector(const DetectorXConfig& cfg): cfg_(cfg) {}
    ~YOLOXDetector(){ cleanup(); }

    bool init(){
        using namespace RivsInference;
        if (inited_) return true;
        LOG("RivsInference init start");
        RivsInference::rivsInference_init();
        handler_ = RivsInference::set_device(cfg_.device_id, cfg_.cluster_ids.data(), cfg_.cluster_ids.size());
        LOG("Device set: device_id=" << cfg_.device_id << ", clusters=" << cfg_.cluster_ids.size());
        std::string exec_path = cfg_.onnx_path + ".exec";

        if (std::filesystem::exists(exec_path)) {
            LOG("Load existing executable: " << exec_path);
            engine_ = RivsInference::create_engine();
            engine_->loadExecutable(exec_path.c_str());
        } else if (std::filesystem::exists(cfg_.onnx_path)) {
            LOG("Build from ONNX: " << cfg_.onnx_path);
            auto* parser = RivsInference::create_parser(RivsInference::TIF_ONNX);
            parser->setInputNames("images");
            parser->setInputShapes("1,3,640,640");
            auto* network = parser->readModel(cfg_.onnx_path.c_str());
            auto* optimizer = RivsInference::create_optimizer();
            optimizer->getConfig()->setBuildFlag(RivsInference::BuildFlag::TIF_KTYPE_MIX_FP16);
            engine_ = optimizer->build(network);
            engine_->saveExecutable(exec_path.c_str());
            RivsInference::release_network(network);
            RivsInference::release_optimizer(optimizer);
            RivsInference::release_parser(parser);
        } else {
            std::cerr << "ONNX not found: " << cfg_.onnx_path << std::endl;
            return false;
        }
        if (!engine_) { LOG("Engine creation failed"); return false; }
        LOG("Engine ready");
        inited_ = true;
        build_grids_strides();
        return true;
    }

    std::vector<std::array<float,6>> detect(const cv::Mat& bgr){
        std::vector<std::array<float,6>> out;
        if (!inited_ && !init()) return out;
        cv::Mat input = letterbox(bgr);
        // HWC BGR -> CHW RGB, keep 0..255 float (no /255 per ref)
        input.convertTo(input, CV_32F);
        cv::cvtColor(input, input, cv::COLOR_BGR2RGB);
        std::vector<float> chw(3 * cfg_.input_w * cfg_.input_h);
        int plane = cfg_.input_w * cfg_.input_h;
        std::vector<cv::Mat> ch = {
            cv::Mat(cfg_.input_h, cfg_.input_w, CV_32F, chw.data()+0*plane),
            cv::Mat(cfg_.input_h, cfg_.input_w, CV_32F, chw.data()+1*plane),
            cv::Mat(cfg_.input_h, cfg_.input_w, CV_32F, chw.data()+2*plane)
        };
        cv::split(input, ch);

        // run
        int out_num = engine_->getOutputNum();
        if (out_num < 1) return out;
        auto odims = engine_->getOutputShape(0);
        int64_t out_elems = 1;
        for (int i=0;i<odims.nbDims;i++) out_elems *= odims.dimension[i];
        std::vector<float> output(out_elems);

        void* input_ptrs[1] = {chw.data()};
        void* output_ptrs[1] = {output.data()};
        bool ok = engine_->runWithBatch(1, input_ptrs, output_ptrs, RivsInference::BufferType::TIF_ENGINE_RSC_IN_HOST_OUT_HOST);
        if (!ok) return out;

        // expected [1, n_preds, 85]
        int n = (odims.nbDims>=2) ? odims.dimension[1] : 0;
        int m = (odims.nbDims>=3) ? odims.dimension[2] : 0;
        const float* det = output.data();
        if (m < 85) return out;

        // decode
        std::vector<cv::Rect> boxes;
        std::vector<float> scores;
        std::vector<int> class_ids;
        boxes.reserve(n);
        scores.reserve(n);
        class_ids.reserve(n);

        for (int i=0;i<n;i++){
            const float* d = det + i*m;
            float raw_x = d[0], raw_y = d[1], raw_w = d[2], raw_h = d[3];
            float obj = d[4];
            float bests=0.f; int best=-1;
            for (int c=5;c<m;c++){
                float v = d[c];
                if (cfg_.apply_sigmoid) v = sigmoid(v);
                if (v > bests){ bests = v; best = c-5; }
            }
            float obj_s = cfg_.apply_sigmoid ? sigmoid(obj) : obj;
            float conf = obj_s * bests;
            if (conf < cfg_.conf_thres) continue;

            // grid-stride decode
            float gx = grids_[i].x;
            float gy = grids_[i].y;
            float st = strides_[i];
            float cx = (raw_x + gx) * st;
            float cy = (raw_y + gy) * st;
            float ww = std::exp(raw_w) * st;
            float hh = std::exp(raw_h) * st;
            int x = int(cx - ww/2);
            int y = int(cy - hh/2);
            boxes.emplace_back(x,y,int(ww),int(hh));
            scores.push_back(conf);
            class_ids.push_back(best);
        }

        if (boxes.empty()) return out;
        std::vector<int> indices;
        cv::dnn::NMSBoxes(boxes, scores, cfg_.conf_thres, cfg_.iou_thres, indices);

        // rescale back by dividing scale only (top-left padding)
        float scale = last_scale_; float pad_x = 0.f; float pad_y = 0.f; // top-left placement
        for (int idx: indices){
            cv::Rect b = boxes[idx];
            float x1 = (b.x - pad_x)/scale;
            float y1 = (b.y - pad_y)/scale;
            float x2 = (b.x + b.width - pad_x)/scale;
            float y2 = (b.y + b.height - pad_y)/scale;
            out.push_back({x1,y1,x2,y2,scores[idx], float(class_ids[idx])});
        }
        return out;
    }

private:
    struct Grid { float x; float y; };
    void build_grids_strides(){
        std::vector<int> strides = cfg_.p6 ? std::vector<int>{8,16,32,64} : std::vector<int>{8,16,32};
        for (int s: strides){
            int hs = cfg_.input_h / s;
            int ws = cfg_.input_w / s;
            for (int y=0;y<hs;y++){
                for (int x=0;x<ws;x++){
                    grids_.push_back({float(x), float(y)});
                    strides_.push_back(float(s));
                }
            }
        }
    }

    cv::Mat letterbox(const cv::Mat& img){
        float r = std::min(float(cfg_.input_w)/img.cols, float(cfg_.input_h)/img.rows);
        int new_w = int(std::round(img.cols*r));
        int new_h = int(std::round(img.rows*r));
        cv::Mat resized; cv::resize(img, resized, cv::Size(new_w,new_h));
        last_scale_ = r; // top-left placement; pad only on right and bottom
        cv::Mat out(cfg_.input_h, cfg_.input_w, CV_8UC3, cv::Scalar(114,114,114));
        resized.copyTo(out(cv::Rect(0, 0, new_w, new_h)));
        return out;
    }

    void cleanup(){
        if (engine_) RivsInference::release_engine(engine_);
        if (handler_) RivsInference::release_device(handler_);
        RivsInference::rivsInference_finish();
        engine_=nullptr; handler_=nullptr; inited_=false;
    }

    DetectorXConfig cfg_;
    RivsInference::IEngine* engine_{nullptr};
    RivsInference::handler_t handler_{};
    bool inited_{false};
    float last_scale_{1.f};
    std::vector<Grid> grids_;
    std::vector<float> strides_;
};

int main(int argc, char** argv){
    std::string rtsp = argc>1 ? argv[1] : otl::replaceHomeDirectory("~/test_data/video/test_car_person_1080P.mp4");
    std::string model_dir = argc>2 ? argv[2] : otl::replaceHomeDirectory("~/models");
    std::string onnx = argc>3 ? argv[3] : "yolox_l.onnx";
    int port = argc>4 ? std::atoi(argv[4]) : 5000;

    LOG("Args: rtsp=" << rtsp << ", model_dir=" << model_dir << ", onnx=" << onnx << ", port=" << port);

    DetectorXConfig cfg; cfg.onnx_path = (std::filesystem::path(model_dir)/onnx).string();
    YOLOXDetector det(cfg);

    cv::VideoCapture cap(rtsp);
    if (!cap.isOpened()){
        std::cerr << "Failed to open: " << rtsp << std::endl; return 1;
    }
    LOG("VideoCapture opened");
    // Prefer FFmpeg for file sources with unknown frame count (e.g., raw H264)
    if (!(rtsp.rfind("rtsp://", 0) == 0) && !(rtsp.rfind("http://", 0) == 0) && !(rtsp.rfind("https://", 0) == 0)) {
        double initial_total = cap.get(cv::CAP_PROP_FRAME_COUNT);
        if (initial_total <= 0) {
            LOG("File source reports invalid frame count (" << initial_total << ") -> reopen with CAP_FFMPEG");
            cap.release();
            std::this_thread::sleep_for(std::chrono::milliseconds(50));
#ifdef CV_VERSION
            cv::VideoCapture cap_ff(rtsp, cv::CAP_FFMPEG);
            if (cap_ff.isOpened()) {
                LOG("Opened with CAP_FFMPEG");
                cap = std::move(cap_ff);
            } else {
                LOG("Failed to open with CAP_FFMPEG, fallback to default backend");
                cap.open(rtsp);
            }
#else
            cap.open(rtsp);
#endif
        }
    }

    std::atomic<bool> running{true};
    WebStatus status; status.last_update = now_time_string();
    std::mutex m; cv::Mat last_frame; std::vector<std::array<float,6>> last_det;
    const bool is_file = !(rtsp.rfind("rtsp://", 0) == 0) && !(rtsp.rfind("http://", 0) == 0) && !(rtsp.rfind("https://", 0) == 0);
    LOG(std::string("Input type: ") + (is_file ? "file" : "stream"));

    // Start web server before detection so we can wait for client connections
    WebMJPEGServer server(port);
    server.setFrameProvider([&](){
        std::lock_guard<std::mutex> lk(m);
        if (last_frame.empty()) return std::vector<uint8_t>{};
        return encode_jpeg(last_frame, 85);
    });
    server.setStatsProvider([&](){ return render_status_html(status); }, [&](){ return status.last_update; });
    server.setStatsJsonProvider([&](){
        std::ostringstream oss;
        oss << "{"
            << "\"fps\":" << status.fps.load() << ","
            << "\"objects\":" << status.total_objects.load() << ","
            << "\"people\":" << status.person_count.load() << ","
            << "\"vehicles\":" << status.car_count.load()
            << "}";
        return oss.str();
    });
    LOG("Starting WebMJPEGServer on port " << port);
    server.start();
    std::cout << "Web server on http://0.0.0.0:" << port << std::endl;

    std::thread worker([&](){
        LOG("Worker thread start");
        auto last_time = std::chrono::steady_clock::now();
        int count=0;
        int read_fail_streak = 0;
        auto last_wait_log = std::chrono::steady_clock::now();
        while (running){
            // Do not start detection until at least one client is connected
            if (server.getClientCount() <= 0) {
                auto noww = std::chrono::steady_clock::now();
                if (std::chrono::duration_cast<std::chrono::milliseconds>(noww - last_wait_log).count() >= 1000) {
                    LOG("Waiting for client connection... no detection running");
                    last_wait_log = noww;
                }
                std::this_thread::sleep_for(std::chrono::milliseconds(100));
                continue;
            }

            cv::Mat frame;
            auto tr0 = std::chrono::steady_clock::now();
            if (!cap.read(frame)){
                auto tr1 = std::chrono::steady_clock::now();
                read_fail_streak++;
                if (is_file) {
                    double total = cap.get(cv::CAP_PROP_FRAME_COUNT);
                    double pos = cap.get(cv::CAP_PROP_POS_FRAMES);
                    if (total > 0) {
                        LOG("File source: pos=" << pos << "/" << total << ", attempt loop to start by seek");
                        cap.set(cv::CAP_PROP_POS_FRAMES, 0);
                        std::this_thread::sleep_for(std::chrono::milliseconds(30));
                    } else {
                        LOG("File source: unknown total frames (" << total << ") -> reopen to loop");
                        cap.release();
                        std::this_thread::sleep_for(std::chrono::milliseconds(80));
#ifdef CV_VERSION
                        cap.open(rtsp, cv::CAP_FFMPEG);
#else
                        cap.open(rtsp);
#endif
                        if (cap.isOpened()) LOG("Reopen file ok"); else LOG("Reopen file failed");
                    }
                    if (read_fail_streak % 50 == 0) {
                        LOG("File source: persistent read failure -> periodic reopen");
                        cap.release();
                        std::this_thread::sleep_for(std::chrono::milliseconds(120));
#ifdef CV_VERSION
                        cap.open(rtsp, cv::CAP_FFMPEG);
#else
                        cap.open(rtsp);
#endif
                        if (cap.isOpened()) LOG("Periodic reopen ok"); else LOG("Periodic reopen failed");
                    }
                } else {
                    if (read_fail_streak % 20 == 0) {
                        LOG("Stream source: attempt reopen");
                        cap.release();
                        std::this_thread::sleep_for(std::chrono::milliseconds(200));
                        cap.open(rtsp);
                        if (cap.isOpened()) LOG("Reopen ok"); else LOG("Reopen failed");
                    }
                    std::this_thread::sleep_for(std::chrono::milliseconds(50));
                }
                continue;
            }
            read_fail_streak = 0;
            auto tr1 = std::chrono::steady_clock::now();

            auto td0 = std::chrono::steady_clock::now();
            auto dets = det.detect(frame);
            auto td1 = std::chrono::steady_clock::now();
            int persons=0, vehicles=0;
            for (auto& d: dets){
                cv::Scalar color(0,0,255);
                cv::rectangle(frame, cv::Point(int(d[0]),int(d[1])), cv::Point(int(d[2]),int(d[3])), color, 2);
                persons += (int)d[5] == 0; // class 0 = person
                int cid = int(d[5]);
                if (cid==2||cid==3||cid==5||cid==7) vehicles++;
            }
            {
                std::lock_guard<std::mutex> lk(m);
                last_frame = frame;
                last_det = dets;
            }
            count++;
            auto now = std::chrono::steady_clock::now();
            float sec = std::chrono::duration<float>(now - last_time).count();
            if (sec >= 1.0f){
                status.fps.store(count/sec); count=0; last_time = now;
                status.total_objects.store((int)dets.size());
                status.person_count.store(persons);
                status.car_count.store(vehicles);
                status.last_update = now_time_string();
                LOG("Status update: fps=" << status.fps.load() << ", objs=" << dets.size()
                    << ", persons=" << persons << ", vehicles=" << vehicles
                    << ", detect_ms=" << std::chrono::duration_cast<std::chrono::milliseconds>(td1 - td0).count());
            }
        }
        LOG("Worker thread end");
    });

    std::cout << "Press Ctrl+C to stop..." << std::endl;
    worker.join();
    server.stop();
    return 0;
}
